import pandas as pd
import requests
import json

def load_data(file_path):
    # return pd.read_excel(file_path)
    return pd.read_csv(file_path)

# def preprocess_data(df):
#     df1 = df.dropna(subset=['问题'])

#     df2 = df1.dropna(how='all', subset=['问题', '答案'])
#     return df2

def analyze_questions(df):
    Q = []
    for i in range(len(df)):
        if df.iloc[i]['sim'] == 0:
            try:
                question = df.iloc[i]['llm_que']
                Q.append(question)
            except Exception as e:
                print(f'Error at index {i+1}: {e}')
                continue
        else:
            Q.append('')
    return Q

def query_responses(questions):
    url = 'http://172.31.30.3:8777/api/local_doc_qa/local_doc_chat'
    A, SOURCE = [], []
    for i, question in enumerate(questions):
        if question:
            headers = {"content-type": "application/json"}
            data = {
                "user_id": "test001",
                "kb_ids": ["kb9816a2566ebb4b3aad7cf50d39ae02ae"],
                "question": question,
                "history": [],
                "rerank": "true",
                "streaming": False,
                "recommend_que": False
            }

            response = requests.request("POST", url, json=data, headers=headers).json()
            if response['code'] == 200:         
                res = response['response']
                print(f'{i+1}条:', res)
                A.append(res)
            else:
                print(('请求失败，状态码:', response.status_code))
        else:
            A.append('')
    return A

def save_result(df, file_path):
    df.to_csv(file_path, index=False, encoding='utf8')

def main():
    test_path = r'/workspace/qanything_local/qanything_kernel/test_data/居住证0416.csv'
    # test_path = r'/workspace/qanything_local/qanything_kernel/test_data/test0305new.csv'
    test_result_path = r'/workspace/qanything_local/qanything_kernel/test_data/new_2_居住证0417.csv'

    df = load_data(test_path)
    
    df = df[500:]
    # df_processed = preprocess_data(df)

    Q = analyze_questions(df)

    A = query_responses(Q)

    df['LLM生成的答案'] = A

    save_result(df, test_result_path)

if __name__ == "__main__":
    main()